# For spatial operations onr aster/vectors
library(terra)

# For plotting operations
library(tidyterra) 
library(ggplot2)
library(mapview)
library(tmap)

# For Perceptually Uniform Colour palettes
library(cetcolor)
library(scico)
 
# Import SMAP soil moisture raster from the downloaded folder
sm=rast("./SampleData-master/raster_files/SMAP_SM.tif") 

print(sm)

# Try:
# dim(sm)   # Dimension (nrow, ncol, nlyr) of the raster
# res(sm)   # X-Y resolution of the raster
# ext(sm)   # Spatial extent of the raster
# crs(sm)   # Coordinate reference system

# Raster plot
terra::plot(sm, main = "Soil Moisture")

library(tidyterra)
library(ggplot2)

ggplot() +
  geom_spatraster(data = sm) +
  scale_fill_gradientn(colors=mypal2,                               # Use user-defined colormap
                       name = "SM",                                 # Colorbar name
                       na.value = "transparent",                    # transparent NA cells
                       labels=(c("0", "0.2", "0.4", "0.6", "0.8")), # Colorbar labels
                       breaks=seq(0,0.8,by=0.2),                    # Set breaks of colorbar
                       limits=c(0,0.8))+                            # Z-axis limits
                       theme_void()                                 # Try different themes: theme_bw(), theme_gray(), theme_minimal

sm_conus= ggplot() +
  geom_spatraster(data = sm) +
  scale_fill_gradientn(colors=mypal2,                               # Use user-defined colormap
                       name = "SM",                                 # Name of the colorbar
                       na.value = "transparent",                    # transparent NA cells
                       labels=(c("0", "0.2", "0.4", "0.6", "0.8")), # Labels of colorbar
                       breaks=seq(0,0.8,by=0.2),                    # Set breaks of colorbar
                       limits=c(0,0.8)) +
  coord_sf(xlim = c(-125,-67),                                      # Add extent for CONUS
              ylim = c(24,50))+               
  theme_bw()                                                        # Try black-and-white theme. 

print(sm_conus)

#We will manipulate the existing tmap_SM plot by adding style parameter:
tmap_SM = tm_shape(sm)+
  tm_raster(alpha = 1, 
            style = "pretty",  # Try other styles: "sd","equal","quantile","kmeans"
            palette = mypal2,  title = "Soil Moisture") +
  tm_layout(legend.position = c("left","bottom"),  # Place legend on bottom left
            inner.margins = 0)+                    # Adjust the legend position
  tm_grid(alpha = 0.2)+                            # Add grid with transparency=0.2
  tm_xlab("Longitude")+                            # X axis label 
  tm_ylab("Latitude")                              # Y axis label 

print(tmap_SM)

# Using breaks would give more control on scale discretization
tmap_SM = tm_shape(sm)+
  tm_raster(alpha = 1, 
            breaks = c(0, 0.1, 0.20, 0.3,0.5, 0.7,1),         # Add color break limits 
            palette = mypal2,                           # Color pal
            title = "Soil Moisture") +
  tm_layout(legend.outside = TRUE,                      # Plot legend outside the box?
            legend.outside.position = c("right","top"), # Legend on top right
            inner.margins = 0)+
  tm_grid(alpha = 0.2)+                                 # Add grid with transparency=0.2
  tm_xlab("Longitude")+                                 # X axis label 
  tm_ylab("Latitude")                                   # Y axis label 
print(tmap_SM)

# Saving the plot to disk as PNG:
tmap_save(tmap_SM,                    # tmap object
          filename = "sm_world.png",  # filename including extension
          dpi = 600,                  # dots per inch 
          height = 5,                 # Height of the exported image
          width = 11,                  # width of the exported image
          units = "in")               # units for width and height

###############################################################
#~~~ PART 1.4.1: Importing and visualizing shapefiles
library(sf)  

# Import the shapefile of global IPCC climate reference regions (only for land) 
IPCC_shp = read_sf("./SampleData-master/CMIP_land/CMIP_land.shp")

# View attribute table of the shapefile
IPCC_shp # Notice the attributes look like a data frame

# Load global coastline shapefile 
coastlines = read_sf("./SampleData-master/ne_10m_coastline/ne_10m_coastline.shp")

# Alternatively, download global coastlines from the web 
# NOTE: May not work if the online server is down
# download.file("https://www.naturalearthdata.com/http//www.naturalearthdata.com/download/110m/physical/ne_110m_coastline.zip?version=4.0.1",destfile = 'ne_110m_coastline.zip')

# # Unzip the downloaded file
# unzip(zipfile = "ne_110m_coastline.zip",exdir = 'ne-coastlines-110m')

# Plot both sf objects using tmap:
tm_shape(IPCC_shp)+                                # Load 'IPCC_shp' in tmap
  tm_borders(col="blue", lwd = 1, lty = "solid")+  # Add IPCC regions as a border in blue
  tm_shape(coastlines)+                            # Initialize 'coastlines' in tmap
  tm_sf()                                          # Add global coastline as sf object

# Subset shapefile for Eastern North-America
terra::plot(IPCC_shp[4,][1], main="Polygon for Eastern North-America")

# Combine tmap plots with overlaying shapefiles
tmap_SM + 
  tm_shape(IPCC_shp)+
  tm_borders()+            
  tm_shape(coastlines)+
  tm_sf()

###############################################################
#~~~ PART 1.4.2: Add spatial point to shapefile/ raster

#~~ Create a spatial point for College Station, Texas
college_station = st_sfc(
                  st_point(x = c(-96.33, 30.62), dim = "XY"), # Lat-long as spatial points 
                  crs = "EPSG:4326")  # Coordinate system: More details in the next part 

#~~ Create map by adding all the layers
tm_shape(IPCC_shp[c(3,4,6,7),])+                       # Selected regions from 'IPCC_shp'
  tm_borders(col = "black",lwd = 1, lty = "solid")+    # Border color
  tm_fill(col = "lightgrey")+                          # Fill color
  tm_shape(coastlines)+                                # Add coastline
  tm_sf(col = "maroon")+                               # Change color of coastline
  tm_shape(college_station)+                           # Add spatial point to the map
  tm_dots(size = 2, col = "blue")                      # Customize point

# Importing SMAP soil moisture data
sm=rast("./SampleData-master/raster_files/SMAP_SM.tif") 

#~~ Projection 1: NAD83 (EPSG: 4269)
sm_proj1 = terra::project(sm, "epsg:4269")

terra::plot(sm_proj1, 
             main = "NAD83",    # Title of the plot
             col = mypal2,      # Colormap for the plot
             axes = FALSE,      # Disable axes
             box = FALSE,       # Disable box around the plots
             asp = NA,          # No fixed aspect ratio; asp=NA fills plot to window
             legend=FALSE)      # Disable legend

#~~ Projection 2: World Robinson projection (ESRI:54030)
sm_proj1 = terra::project(sm, "ESRI:54030")
terra::plot(sm_proj1, 
             main = "Robinson", # Title of the plot
             col = mypal2,      # Colormap for the plot
             axes = FALSE,      # Disable axes
             box = FALSE,       # Disable box around the plots
             asp = NA,          # No fixed aspect ratio; asp=NA fills plot to window
             legend=FALSE)      # Disable legend

#~~~Projection 2: Robinson projection

WorldSHP=terra::vect(spData::world)

RobinsonPlot <- ggplot() +
  geom_spatraster(data = sm)+                   # Plot SpatRaster layer               
  geom_spatvector(data = WorldSHP, 
                  fill = "transparent") +       # Add world political map
  ggtitle("Robinson Projection") +              # Add title
  scale_fill_gradientn(colors=mypal2,           # Use user-defined colormap
                       name = "Soil Moisture",  # Name of the colorbar
                       na.value = "transparent",# Set color for NA values
                       lim=c(0,0.8))+           # Z axis limit
  theme_minimal()+                              # Select theme. Try 'theme_void'
  theme(plot.title = element_text(hjust =0.5),  # Place title in the middle of the plot
        text = element_text(size = 12))+        # Adjust plot text size for visibility
  coord_sf(crs = "ESRI:54030",                  # Reproject to World Robinson
           xlim = c(-152,152)*100000,    
           ylim = c(-55,90)*100000)

print(RobinsonPlot)

WorldSHP = st_as_sf(WorldSHP)           # Convert 'WorldSHP' to simple feature

tm_shape(WorldSHP,                      # Initiate shapefile       
         projection = 'ESRI:54030',     # Set projection: World Robinson
         ylim = c(-65, 90)*100000,      # Set y-limit
         xlim = c(-152,152)*100000,     # Set x-limit
         raster.warp = TRUE)+              
  tm_sf()+                              # Plot shapefile 
  tm_shape(sm,                          # Add raster file
           projection = 'ESRI:54030',   # Set projection 
           raster.warp = FALSE) +
  tm_raster( palette = mypal2,          # Set color map for raster
             title = "Soil Moisture")+  # Add plot title
  tm_layout(main.title = "Surface Soil Moisture",
            main.title.fontfamily = "Times",               # Set text font
            legend.show = T,                               # Show legend= T/F
            legend.outside = T,                            
            legend.outside.position = c("right", "top"),   # Legend position
            frame = FALSE,                                 # Add plot frame
            earth.boundary = c(-160, -65, 160, 88),        # Set plot boundary
            earth.boundary.color = "grey",                 # Boundary color
            earth.boundary.lwd = 2,                        # Boundary linewidth
            fontfamily = "Times")+                         # Text font
  tm_graticules(alpha = 0.6,                               # Add lat-long graticules
                col = "lightgrey")

# Original resoluton of raster for reference
res(sm)

#~~ Aggregate raster to coarser resolution
SMcoarse = terra::aggregate(sm,           # Soil moisture raster
                            fact = 10,    # Aggregate by x 10
                            fun = mean)   # Function used to aggregate values
res(SMcoarse)

#~~ Disaggregate raster to finer resolution
SMfine = terra::disagg(sm, 
                   fact=3, 
                   method='bilinear')
res(SMfine)

#~~ Raster resampling
# Import global aridity raster
aridity=rast("./SampleData-master/raster_files/aridity_36km.tif") 
# Plot aridity map
terra::plot(aridity, col=mypal2)

# Resample aridity raster to coarse resolution  
aridityResamp=terra::resample(aridity,      # Original raster
                       SMcoarse,     # Target resolution raster
                       method='ngb') # bilinear or ngb (nearest neighbor) 
# Plot resampled aridity map
terra::plot(aridityResamp, col=mypal2)

# Simple arithmetic operations
sm2=sm*2
print(sm2) # Try sm2=sm*10, or sm2=sm^2 and see the difference in sm2 values

# Summary statistics
global(sm, mean, na.rm = T)

global(sm, sd, na.rm = T)

global(sm, quantile, probs = c(0.25, 0.75), na.rm = T)

# User-defined statistics by defining own function
quant_fun = function(x){ 
  quantile(x, probs = c(0.25, 0.75), na.rm=TRUE)
} 
global(sm, quant_fun)   # 25th, and 75th percentile of each layer

# Note: With a multi-layered raster object, global() will summarize each layer separately.

###############################################################
#~~~ PART 2.3.1: Using shapefile to summarize a raster
sm_IPCC_df=terra::extract(sm,        # Raster to be summarized
                          IPCC_shp,  # Shapefile/ polygon to summarize the raster
                          df=TRUE,   # Gives the summary statistics as a dataframe
                          fun=mean,  # Desired statistic: mean, sum, min and max 
                          na.rm=TRUE)# Ignore NA values? TRUE=yes! 

head(sm_IPCC_df)

###############################################################
#~~~ PART 2.3.2: Extract cell values for each region 
sm_IPCC_list=terra::extract(sm,       # Raster to be summarized
                          IPCC_shp,   # Shapefile/ polygon to summarize the raster
                          df=FALSE,   # Returns a list
                          fun=NULL,   # fun=NULL will output cell values within each region
                          na.rm=TRUE) # Ignore NA values? yes! 

# Apply function on cell values for each region
library(tidyverse)

sm_IPCC_list %>% 
  as_tibble() %>% 
  group_by(ID) %>% 
  summarise(mean_SM = mean(SMAP_SM, na.rm =T))

#~~ Try user defined function
myfun=function (y){return(mean(y, na.rm=TRUE))}    # User defined function for calculating means

sm_IPCC_list %>% 
  as_tibble() %>% 
  group_by(ID) %>% 
  summarise(mean_SM = myfun(SMAP_SM)) 

# Is this the same as the previous result?

###############################################################
#~~~ PART 2.2.3: Convert a raster to a shapefile
aridity=rast("./SampleData-master/raster_files/aridity_36km.tif") #Global aridity

# Convert raster to shapefile
arid_poly=st_as_sf(as.polygons(aridity))   # Convert SpatRaster to polygon and then to sf

# Plot aridity polygon
plot(arid_poly, 
     col=arid_poly$aridity_36km)  # Colors based on aridity values (i.e. 1,2,3,4,5)

sm_arid_df=terra::extract(sm,        # Raster to be summarized
                          arid_poly, # Shapefile/ polygon to summarize the raster
                          df=TRUE,   # Gives the summary statistics as a dataframe
                          fun=mean,  # Desired statistic: mean, sum, min and max 
                          na.rm=TRUE)# Ignore NA values? yes! 

# Lets plot the climate-wise mean of surface soil moisture
plot(sm_arid_df,     
     xaxt = "n",              # Disable x-tick labels
     xlab="Aridity",          # X axis label
     ylab="Soil moisture",    # Y axis label
     type="b",                # line type
     col="blue",              # Line color
     main="Climate-wise mean of surface soil moisture")
axis(1, at=1:5, labels=c("Hyper-arid", "Arid", "Semi-Arid","Sub-humid","Humid"))

###############################################################
#~~~ PART 3.1.1: Create and plot NDVI SpatRaster
library(terra)
# Location of the NDVI raster files
ndvi_path="./SampleData-master/raster_files/NDVI/" #Specify location of NDVI rasters

# List of all NDVI rasters
ras_path=list.files(ndvi_path,pattern='*.tif',full.names=TRUE)
head(ras_path)

# Method 1: Use lapply to create raster layer list from the raster location
ras_list = lapply(paste(ras_path, sep = ""), rast)
#This a list of 23 raster objects stored as individual elements.

#Convert raster layer lists to data cube 
ras_stack = rast(ras_list)     # Stacking all rasters as a data cube!!! 
# This a multi-layer (23 layers in this case) SpatRaster Object.

#inMemory() reports whether the raster data is stored in memory or on disk.
inMemory(ras_stack[[1]]) 

# Method 2: Using pipes to create raster layers from the raster location
library(tidyr) # For piping
ras_list = ras_path %>%  purrr::map(~ rast(.x))  # Import the raster
ras_stack = rast(ras_list)  # Convert RasterStack to RasterBrick

# Method 3: Using for loop to create raster layers from the raster location
ras_stack=rast()
for (nRun in 1:length(ras_path)){
  ras_stack=c(ras_stack,rast(ras_path[[nRun]]))
}

# Check dimension of data cube
dim(ras_stack) #[x: y: z]- 23 raster layers with 456 x 964 cells

# Number of layers in raster brick
nlyr(ras_stack)

# Check variable names 
names(ras_stack)

# Subset raster stack/brick (notice the double [[]] bracket and similarity to lists)
sub_ras_stack=ras_stack[[c(1,3,5,10,12)]] #Select 1st, 3rd, 5th, 10th and 12th layers

#Try subsetting with dates:
Season<-str_subset(string = names(ras_stack), pattern = c("20..-0[3/4/5]")) 
ras_stack[[Season]]

#~~ Plot first 4 elements of NDVI SpatRaster
# Function to add shapefile to a raster plot
addCoastlines=function(){
  plot(vect(coastlines), add=TRUE)   #convert 'coastlines' to vector format
  } 

terra::plot(sub_ras_stack[[1:4]],  # Select raster layers to plot
     col=mypal2,                   # Specify colormap
     asp=NA,                       # Aspect ratio= NA
     fun=addCoastlines)            # Add coastline to each layer

###############################################################
#~~~ PART 3.3.1: Layer-wise cell-statistics 
# Layer-wise Mean
global(sub_ras_stack, mean, na.rm= T) # Mean of each raster layer. Try modal, median, min etc. 

# Layer-wise quantiles
global(sub_ras_stack, quantile, probs=c(.25,.75), na.rm= T)

# User-defined statistics by defining own function
quant_fun = function(x) {quantile(x, probs = c(0.25, 0.75), na.rm=TRUE)} 
global(sub_ras_stack, quant_fun) # 25th, and 75th percentile of each layer

# Custom function for mean, variance and skewness
my_fun = function(x){ 
  meanVal=mean(x, na.rm=TRUE)              # Mean 
  varVal=var(x, na.rm=TRUE)                # Variance
  skewVal=moments::skewness(x, na.rm=TRUE) # Skewness
  output=c(meanVal,varVal,skewVal)         # Combine all statistics
  names(output)=c("Mean", "Var","Skew")    # Rename output variables
  return(output)                           # Return output
} 

global(sub_ras_stack, my_fun) # Mean, Variance and skewness of each layer

#You can also use summary() to retrieve common descriptive statistics for each layer
summary(sub_ras_stack)

###############################################################
#~~~ PART 3.3.2: Global cell-statistics

# By vectorizing the SpatRaster and finding statistics
min_val=min(as.vector(sub_ras_stack), na.rm=T) 
max_val=max(as.vector(sub_ras_stack), na.rm=T)
print(c(min_val,max_val))

# Another example of statistics of vectorized SpatRaster
quant=quantile(as.vector(sub_ras_stack), c(0.01,0.99), na.rm=TRUE) #1st and 99th percentiles
print(quant)

###############################################################
#~~~ PART 3.3.3: Layer-wise operations on SpatRaster

# Arithmetic operations on SpatRaster are same as lists
add = sub_ras_stack+10                  # Add a number to raster layers
mult = sub_ras_stack*5                  # Multiply a number to raster layers
subset_mult = sub_ras_stack[[1:3]]*10   # Multiply a number to a subset of raster layers

# Data filtering based on cell-value
filter_stack = sub_ras_stack            # Create a RasterBrick to operate on
filter_stack[filter_stack<0.5] = NA     # Assign NA to any value less than 0.5

# Let's plot the filtered rasters
plot(filter_stack, 
     col=mypal2,  # Color pal  
     asp=NA,      # Aspect ratio: NA, fill to plot space
     nc=2,        # Number of columns to arrange plots
    fun=function(){plot(vect(coastlines), add=TRUE)} # Add background map
)

# Layer-wise Log-transformation
log_ras_stack=log(sub_ras_stack) 

# Normalize raster layers to [0,1] based on min and max of raster brick/stack
norm_stk=(sub_ras_stack-min_val)/(max_val-min_val) # Notice that the values are between [0,1]

# Plot in Robinson projection
mypal3 = cetcolor::cet_pal(20, name = "l5")
WorldSHP = st_as_sf(spData::world)

norm_NDVI= tm_shape(WorldSHP, projection = 'ESRI:54030', ylim = c(-65, 90)*100000, 
         xlim = c(-152,152)*100000, raster.warp = T)+
  tm_sf()+
  tm_shape(norm_stk[[2:5]], projection = 'ESRI:54030', raster.warp = FALSE) +
  tm_raster(palette = mypal3, title = "Soil Moisture", style = "cont")+
  tm_layout(main.title = "NDVI",
            main.title.fontfamily = "Times",
            legend.show = T,
            legend.outside = T,
            legend.outside.position = c("right", "top"),
            frame = FALSE, 
            earth.boundary = c(-160, -65, 160, 88),
            earth.boundary.color = "grey",
            earth.boundary.lwd = 2,
            fontfamily = "Times")+
  tm_graticules(alpha = 0.6, col = "lightgrey")+
  tm_facets(ncol = 2)

print(norm_NDVI)

# 5.3.1. Plot aridity map
tm_shape(aridity)+
  tm_raster(alpha = 1,      # Transparency
            style = "cat",  # cat= categorical 
            title = "Global aridity Regimes") +
  tm_layout(legend.outside = TRUE,
            legend.outside.position = c("right","top"), inner.margins = 0)+
  tm_grid(alpha = 0.2)+ tm_xlab("Longitude")+ tm_ylab("Latitude")

# 5.3.2. Plot resamples aridity maps
tm_shape(aridityResamp)+
  tm_raster(alpha = 1, 
            style = "cat",  # cat= categorical 
            title = "Resampled Global Aridity") +
  tm_layout(legend.outside = TRUE,
            legend.outside.position = c("right","top"), inner.margins = 0)+
  tm_grid(alpha = 0.2)+ 
  tm_xlab("Longitude")+ 
  tm_ylab("Latitude")

# 5.3.3. Plot first 4 elements of NDVI SpatRaster
tm_shape(sub_ras_stack[[1:4]])+
  tm_raster(colorNA = "white",
            showNA = F, 
            breaks = c(0.0, 0.2, 0.4, 0.6, 0.8, 1.0), 
            title = "Global NDVI")+
  tm_shape(coastlines)+
  tm_sf()

# 5.3.4. Plot cropped raster
tm_shape(usa_crop[[1:4]])+
  tm_raster(title = "Cropped CONUS NDVI", style = "cont")+
  tm_grid(alpha = 0.2)+ tm_xlab("Longitude")+ tm_ylab("Latitude")+
  tm_layout(inner.margins = 0)+
  tm_shape(conus)+
  tm_sf(alpha = 0.2)

# 5.3.5. Mask SpatRaster using USA polygon
tm_shape(ndvi_mask_usa[[1:4]])+
  tm_raster(title = "Masked CONUS NDVI", breaks = c(0.0, 0.2, 0.4, 0.6, 0.8))+
  tm_grid(alpha = 0.2)+ tm_xlab("Longitude")+ tm_ylab("Latitude")+
  tm_layout(inner.margins = 0)+
  tm_shape(conus)+
  tm_sf(alpha = 0.2)

# 5.3.6. Let's plot the filtered rasters
tm_shape(filter_stack)+
  tm_raster()+
  tm_shape(coastlines)+                 # Add coastlines to each plot
  tm_sf()       

# 5.3.7. Plot raster and shapefile 
tm_shape(ndvi_mask_states[[1]])+   # '[[1]] selects the first layer for plotting
  tm_raster(breaks = c(0.0, 0.2, 0.4, 0.6, 0.8, 1.0), title = "NDVI")+
  tm_shape(state_plot)+
  tm_polygons(alpha = 0.1)+
  tm_layout(legend.outside = T)+
  tm_grid(alpha = 0.5, col = "lightgrey")

#  5.3.8. Plot statistics
output1<- tm_shape(stat_ras[[1]], projection = 'ESRI:54030', raster.warp = FALSE,
                   ylim = c(-65, 90)*100000, xlim = c(-152,152)*100000)+
  tm_raster(title = "Mean", style = "cont", colorNA = "white", showNA = F)+
  tm_layout(legend.outside = T, legend.outside.position = c("right", "bottom"), 
            frame = F, earth.boundary =  c(-165, -65, 165, 88))+
  tm_shape(coastlines, projection = 'ESRI:54030')+
  tm_sf()+ 
  tm_graticules(alpha = 0.6, col = "lightgrey")

output2<- tm_shape(stat_ras[[2]], projection = 'ESRI:54030', raster.warp = FALSE,
                   ylim = c(-65, 90)*100000, xlim = c(-152,152)*100000)+
  tm_raster(title = "Variance", style = "cont", colorNA = "white", showNA = F)+
  tm_layout(legend.outside = T, legend.outside.position = c("right", "bottom"), 
            frame = F, earth.boundary = c(-165, -65, 165, 88))+
  tm_shape(coastlines, projection = 'ESRI:54030')+
  tm_sf()

  tmap_arrange(output1, output2, ncol = 1)

#Plot NDVI difference
tm_shape(diff_rast)+
  tm_raster(palette = "RdYlGn", title = "Difference in NDVI", legend.hist = T)+ 
  tm_grid(alpha = 0.2)+ 
  tm_xlab("Longitude")+ 
  tm_ylab("Latitude")+
  tm_style("cobalt")+
  tm_layout(legend.position = c("left", "bottom"), legend.hist.size = 2)